{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-38edKrF8ylR"
      },
      "source": [
        "# 官方数据集下载 + 解压\n",
        "（所有外部数据集通过wget导入在本模块导入，在其他部分导入可能导致失去成绩）"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "F5N6wkdR75nJ",
        "outputId": "1df85a0d-5958-4a1d-9a20-a5acd16745da"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--2024-09-05 17:39:53--  https://drive.usercontent.google.com/download?id=1Y0ejShdcx6YiaE3jJ0yxcOTKxirR-MAh&export=download&authuser=0&confirm=t&uuid=90e67168-c423-4396-a7ac-4cfcb1f2cc01&at=APZUnTXTP3YsMMcYVNsvI5-B_UtN:1720979862404\n",
            "Resolving drive.usercontent.google.com (drive.usercontent.google.com)... 172.217.203.132, 2607:f8b0:400c:c07::84\n",
            "Connecting to drive.usercontent.google.com (drive.usercontent.google.com)|172.217.203.132|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1089154958 (1.0G) [application/octet-stream]\n",
            "Saving to: ‘CIKM_Training.zip’\n",
            "\n",
            "CIKM_Training.zip   100%[===================>]   1.01G  67.4MB/s    in 20s     \n",
            "\n",
            "2024-09-05 17:40:15 (52.2 MB/s) - ‘CIKM_Training.zip’ saved [1089154958/1089154958]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "!wget --header=\"Host: drive.usercontent.google.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Cookie: __Secure-ENID=12.SE=Yd0Bj-CLJ14fnd4qzdJHmwUs4B5zz46UaPC1cPJigNqqFV9PtM2CYyBpSbCkOyzUwzlEdZ1nZFf-igtGi7wSdJ_gqQSfQfh84r9egqFQAy9-GKayCRbdQKdera-2mkpuIT-c64CyR9vfNojM3hxZ9Dej-dGvtxlGjal9ttEHybw; __gsas=ID=ae0421b9a34b478c:T=1710758437:RT=1710758437:S=ALNI_MZP13R9ZOHbCzC0rgHSMrGXj6GCsg; SEARCH_SAMESITE=CgQI35oB; SID=g.a000lgiBaVVDJxM_nyJHv8SGKdW8l-Ea7eoYBmsE1puBGqeYmuoUD3CsQuje5kFlDuyAjS6tbwACgYKAewSARQSFQHGX2MiQwM3sBqt9QySR0B5yipflRoVAUF8yKo6yV91BccjS1Z6o32Ux6ew0076; __Secure-1PSID=g.a000lgiBaVVDJxM_nyJHv8SGKdW8l-Ea7eoYBmsE1puBGqeYmuoUdkz-2WaNotDz3XhWHdZaewACgYKAdMSARQSFQHGX2Mi2DyEcyhQnLt_3xtpRFnETBoVAUF8yKoObySYGWwftPorBYI_hdIh0076; __Secure-3PSID=g.a000lgiBaVVDJxM_nyJHv8SGKdW8l-Ea7eoYBmsE1puBGqeYmuoUMoY9B1xAja7hD33yfC_l8gACgYKAXESARQSFQHGX2Miy7Eh3rWIzqh1GS_d-315mhoVAUF8yKpBtafXePsFsQwbeu_1xhEZ0076; HSID=AqK8DRIVZrNtHDQEM; SSID=AF8GbXclnoQKbndEK; APISID=eaYPuocswpo7W9-H/AOCjbMPSZJ096e5h4; SAPISID=xTIl5TE60VtPYf3j/A-54gZvfpsOMhgwjL; __Secure-1PAPISID=xTIl5TE60VtPYf3j/A-54gZvfpsOMhgwjL; __Secure-3PAPISID=xTIl5TE60VtPYf3j/A-54gZvfpsOMhgwjL; NID=515=plru-gK0AjVf1Sy8SC5LvpAocQ8efRsdx0CZ1kdvR10iJMaeT-fsU6T5AyKgdFpWE-kImnEI5utU-487cxMnVkrz0_-aB-3blUbOI1wu31UI-2dAjxd2LWGEgubs2zJi8cjE8hfCzP-TgVqkwzSNkE6riAOmKvZvkdSWWgnwFiQLxr3X7et1TNvcPZvpEHzG8iqoylaY55lZje6L4CmiT0b8N2PzagLdpOZdUGkWyRJDlTsAeSnlvMM5M8ZTWjZm7UvVNqMRBh-CA4KJ-BkCavo-Pt0F4-3-qSceoZ3nNz6AUfFk_2J9dWLmijh23NbSZsQP4z2u7O2WDMwTmntI1Qc5QiIuXjxGDLDz6GQPSEMIyWZCRTsNxZtmNlKQROEsPwKKMmmylxBjRXLFtSKyLitb6oDPhkp8YbUh984lDVAuZkAOcjAuCgYFdAEmNyGGrWWFPmdaUAUZ-uxIL8lClSnbNx0zECu88QhO-QtallWvKzDYr3ZK7czaeJi2nQX53eUWfJM0eyOeeric1Ltf1wA1_w0zXPaBF5rPBONe7Rzm5phafNrSAl9XpofCiBT_B4mmb81hfFlGrBI3_4Rtj7UohFGwMaKcrIyPLBok4nQogBBPAzhPsnNFLJgBOUqQLAMDIoup982yf7v8ISeApMkZdgQ; __Secure-1PSIDTS=sidts-CjEB4E2dkZiFPZveD2L67T0j-mwN6ErulV1FpkkkhAlfHtcpBwyaLEFAl9KPDrhmId99EAA; __Secure-3PSIDTS=sidts-CjEB4E2dkZiFPZveD2L67T0j-mwN6ErulV1FpkkkhAlfHtcpBwyaLEFAl9KPDrhmId99EAA; SIDCC=AKEyXzURbTd6qCpA02QoGXCORelVic8I05vY0Fudj1oUukq6szKRAuO4mcverQ_TFNKg0pPsF-tiZg; __Secure-1PSIDCC=AKEyXzVPd347dt-7_3JSjQc3gmzPqwRtmRgaQpuVqN0qtN34VORpmjxOBVKFZsXxU0m0JibWfONm; __Secure-3PSIDCC=AKEyXzWa_JoS5Duud9ECYjQhLb8ENX9_ML6X5YK4El1eSaET9Cq21vXwSW53TCFtQ0e-KwYt3BEK\" --header=\"Connection: keep-alive\" \"https://drive.usercontent.google.com/download?id=1Y0ejShdcx6YiaE3jJ0yxcOTKxirR-MAh&export=download&authuser=0&confirm=t&uuid=90e67168-c423-4396-a7ac-4cfcb1f2cc01&at=APZUnTXTP3YsMMcYVNsvI5-B_UtN:1720979862404\" -c -O 'CIKM_Training.zip'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_T7GT7nZZYHG",
        "outputId": "0bc413db-01f2-4b6a-c641-c382b16bbe1b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--2024-09-05 17:40:15--  https://drive.usercontent.google.com/download?id=17M0nVsBdeiZxMnFVhKMuq89APc9pbA91&export=download&authuser=0&confirm=t&uuid=de518810-0b86-44da-9924-b3044aedae65&at=APZUnTWx_R5StqP3iGT6U7FBhIm-:1720979898499\n",
            "Resolving drive.usercontent.google.com (drive.usercontent.google.com)... 172.217.203.132, 2607:f8b0:400c:c07::84\n",
            "Connecting to drive.usercontent.google.com (drive.usercontent.google.com)|172.217.203.132|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 94200676 (90M) [application/octet-stream]\n",
            "Saving to: ‘CIKM_Test.zip’\n",
            "\n",
            "CIKM_Test.zip       100%[===================>]  89.84M  36.7MB/s    in 2.4s    \n",
            "\n",
            "2024-09-05 17:40:19 (36.7 MB/s) - ‘CIKM_Test.zip’ saved [94200676/94200676]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "!wget --header=\"Host: drive.usercontent.google.com\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7\" --header=\"Accept-Language: zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6\" --header=\"Cookie: __Secure-ENID=12.SE=Yd0Bj-CLJ14fnd4qzdJHmwUs4B5zz46UaPC1cPJigNqqFV9PtM2CYyBpSbCkOyzUwzlEdZ1nZFf-igtGi7wSdJ_gqQSfQfh84r9egqFQAy9-GKayCRbdQKdera-2mkpuIT-c64CyR9vfNojM3hxZ9Dej-dGvtxlGjal9ttEHybw; __gsas=ID=ae0421b9a34b478c:T=1710758437:RT=1710758437:S=ALNI_MZP13R9ZOHbCzC0rgHSMrGXj6GCsg; SEARCH_SAMESITE=CgQI35oB; SID=g.a000lgiBaVVDJxM_nyJHv8SGKdW8l-Ea7eoYBmsE1puBGqeYmuoUD3CsQuje5kFlDuyAjS6tbwACgYKAewSARQSFQHGX2MiQwM3sBqt9QySR0B5yipflRoVAUF8yKo6yV91BccjS1Z6o32Ux6ew0076; __Secure-1PSID=g.a000lgiBaVVDJxM_nyJHv8SGKdW8l-Ea7eoYBmsE1puBGqeYmuoUdkz-2WaNotDz3XhWHdZaewACgYKAdMSARQSFQHGX2Mi2DyEcyhQnLt_3xtpRFnETBoVAUF8yKoObySYGWwftPorBYI_hdIh0076; __Secure-3PSID=g.a000lgiBaVVDJxM_nyJHv8SGKdW8l-Ea7eoYBmsE1puBGqeYmuoUMoY9B1xAja7hD33yfC_l8gACgYKAXESARQSFQHGX2Miy7Eh3rWIzqh1GS_d-315mhoVAUF8yKpBtafXePsFsQwbeu_1xhEZ0076; HSID=AqK8DRIVZrNtHDQEM; SSID=AF8GbXclnoQKbndEK; APISID=eaYPuocswpo7W9-H/AOCjbMPSZJ096e5h4; SAPISID=xTIl5TE60VtPYf3j/A-54gZvfpsOMhgwjL; __Secure-1PAPISID=xTIl5TE60VtPYf3j/A-54gZvfpsOMhgwjL; __Secure-3PAPISID=xTIl5TE60VtPYf3j/A-54gZvfpsOMhgwjL; NID=515=plru-gK0AjVf1Sy8SC5LvpAocQ8efRsdx0CZ1kdvR10iJMaeT-fsU6T5AyKgdFpWE-kImnEI5utU-487cxMnVkrz0_-aB-3blUbOI1wu31UI-2dAjxd2LWGEgubs2zJi8cjE8hfCzP-TgVqkwzSNkE6riAOmKvZvkdSWWgnwFiQLxr3X7et1TNvcPZvpEHzG8iqoylaY55lZje6L4CmiT0b8N2PzagLdpOZdUGkWyRJDlTsAeSnlvMM5M8ZTWjZm7UvVNqMRBh-CA4KJ-BkCavo-Pt0F4-3-qSceoZ3nNz6AUfFk_2J9dWLmijh23NbSZsQP4z2u7O2WDMwTmntI1Qc5QiIuXjxGDLDz6GQPSEMIyWZCRTsNxZtmNlKQROEsPwKKMmmylxBjRXLFtSKyLitb6oDPhkp8YbUh984lDVAuZkAOcjAuCgYFdAEmNyGGrWWFPmdaUAUZ-uxIL8lClSnbNx0zECu88QhO-QtallWvKzDYr3ZK7czaeJi2nQX53eUWfJM0eyOeeric1Ltf1wA1_w0zXPaBF5rPBONe7Rzm5phafNrSAl9XpofCiBT_B4mmb81hfFlGrBI3_4Rtj7UohFGwMaKcrIyPLBok4nQogBBPAzhPsnNFLJgBOUqQLAMDIoup982yf7v8ISeApMkZdgQ; __Secure-1PSIDTS=sidts-CjEB4E2dkZiFPZveD2L67T0j-mwN6ErulV1FpkkkhAlfHtcpBwyaLEFAl9KPDrhmId99EAA; __Secure-3PSIDTS=sidts-CjEB4E2dkZiFPZveD2L67T0j-mwN6ErulV1FpkkkhAlfHtcpBwyaLEFAl9KPDrhmId99EAA; SIDCC=AKEyXzV9L2DyLa7cOj7QEH5jq731oO3bjUfxB9xunhlp0aLuj_Yrks1puUhcLsZ8rELAfcHJjDmdqg; __Secure-1PSIDCC=AKEyXzUqTJxy1COupBctW9fNufvYQ4ZxCVzD4btsqLhg-i8UKiKC1jI3mI_5nliA_a0z-FM7PD0I; __Secure-3PSIDCC=AKEyXzU45-0-sVl11D0Mg7O6UieJpansLMEGJeXrmWGqtoNEtTeyHkgewKyRGMZVmJHm49Bzj1IM\" --header=\"Connection: keep-alive\" \"https://drive.usercontent.google.com/download?id=17M0nVsBdeiZxMnFVhKMuq89APc9pbA91&export=download&authuser=0&confirm=t&uuid=de518810-0b86-44da-9924-b3044aedae65&at=APZUnTWx_R5StqP3iGT6U7FBhIm-:1720979898499\" -c -O 'CIKM_Test.zip'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "aPTwHZWpaJ-1"
      },
      "outputs": [],
      "source": [
        "!unzip -q CIKM_Training.zip\n",
        "!unzip -q CIKM_Test.zip"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# 解压后的数据目录\n",
        "# Training\n",
        "#     Dataset_1\n",
        "#         Feature_File\n",
        "#         Label_File\n",
        "#     Dataset_2\n",
        "#         Feature_File\n",
        "#         Label_File\n",
        "\n",
        "# Test\n",
        "#     Dataset_1\n",
        "#         Feature_File\n",
        "#     Dataset_2\n",
        "#         Feature_File\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XxaSCHvUMm4X"
      },
      "source": [
        "# 常见工具安装与导入\n",
        "（所有外部代码和包需要在本模块导入）"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kD437pLQMvOC",
        "outputId": "9ca52c4e-8d1b-49cd-bf30-98a377e70da1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting open3d\n",
            "  Downloading open3d-0.18.0-cp310-cp310-manylinux_2_27_x86_64.whl.metadata (4.2 kB)\n",
            "Requirement already satisfied: numpy>=1.18.0 in /usr/local/lib/python3.10/dist-packages (from open3d) (1.26.4)\n",
            "Collecting dash>=2.6.0 (from open3d)\n",
            "  Downloading dash-2.18.0-py3-none-any.whl.metadata (10 kB)\n",
            "Requirement already satisfied: werkzeug>=2.2.3 in /usr/local/lib/python3.10/dist-packages (from open3d) (3.0.4)\n",
            "Requirement already satisfied: nbformat>=5.7.0 in /usr/local/lib/python3.10/dist-packages (from open3d) (5.10.4)\n",
            "Collecting configargparse (from open3d)\n",
            "  Downloading ConfigArgParse-1.7-py3-none-any.whl.metadata (23 kB)\n",
            "Collecting ipywidgets>=8.0.4 (from open3d)\n",
            "  Downloading ipywidgets-8.1.5-py3-none-any.whl.metadata (2.3 kB)\n",
            "Collecting addict (from open3d)\n",
            "  Downloading addict-2.4.0-py3-none-any.whl.metadata (1.0 kB)\n",
            "Requirement already satisfied: pillow>=9.3.0 in /usr/local/lib/python3.10/dist-packages (from open3d) (9.4.0)\n",
            "Requirement already satisfied: matplotlib>=3 in /usr/local/lib/python3.10/dist-packages (from open3d) (3.7.1)\n",
            "Requirement already satisfied: pandas>=1.0 in /usr/local/lib/python3.10/dist-packages (from open3d) (2.1.4)\n",
            "Requirement already satisfied: pyyaml>=5.4.1 in /usr/local/lib/python3.10/dist-packages (from open3d) (6.0.2)\n",
            "Requirement already satisfied: scikit-learn>=0.21 in /usr/local/lib/python3.10/dist-packages (from open3d) (1.3.2)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from open3d) (4.66.5)\n",
            "Collecting pyquaternion (from open3d)\n",
            "  Downloading pyquaternion-0.9.9-py3-none-any.whl.metadata (1.4 kB)\n",
            "Requirement already satisfied: Flask<3.1,>=1.0.4 in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (2.2.5)\n",
            "Requirement already satisfied: plotly>=5.0.0 in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (5.15.0)\n",
            "Collecting dash-html-components==2.0.0 (from dash>=2.6.0->open3d)\n",
            "  Downloading dash_html_components-2.0.0-py3-none-any.whl.metadata (3.8 kB)\n",
            "Collecting dash-core-components==2.0.0 (from dash>=2.6.0->open3d)\n",
            "  Downloading dash_core_components-2.0.0-py3-none-any.whl.metadata (2.9 kB)\n",
            "Collecting dash-table==5.0.0 (from dash>=2.6.0->open3d)\n",
            "  Downloading dash_table-5.0.0-py3-none-any.whl.metadata (2.4 kB)\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (8.4.0)\n",
            "Requirement already satisfied: typing-extensions>=4.1.1 in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (4.12.2)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (2.32.3)\n",
            "Collecting retrying (from dash>=2.6.0->open3d)\n",
            "  Downloading retrying-1.3.4-py3-none-any.whl.metadata (6.9 kB)\n",
            "Requirement already satisfied: nest-asyncio in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (1.6.0)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from dash>=2.6.0->open3d) (71.0.4)\n",
            "Collecting comm>=0.1.3 (from ipywidgets>=8.0.4->open3d)\n",
            "  Downloading comm-0.2.2-py3-none-any.whl.metadata (3.7 kB)\n",
            "Requirement already satisfied: ipython>=6.1.0 in /usr/local/lib/python3.10/dist-packages (from ipywidgets>=8.0.4->open3d) (7.34.0)\n",
            "Requirement already satisfied: traitlets>=4.3.1 in /usr/local/lib/python3.10/dist-packages (from ipywidgets>=8.0.4->open3d) (5.7.1)\n",
            "Collecting widgetsnbextension~=4.0.12 (from ipywidgets>=8.0.4->open3d)\n",
            "  Downloading widgetsnbextension-4.0.13-py3-none-any.whl.metadata (1.6 kB)\n",
            "Requirement already satisfied: jupyterlab-widgets~=3.0.12 in /usr/local/lib/python3.10/dist-packages (from ipywidgets>=8.0.4->open3d) (3.0.13)\n",
            "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3->open3d) (1.3.0)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3->open3d) (0.12.1)\n",
            "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3->open3d) (4.53.1)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3->open3d) (1.4.5)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3->open3d) (24.1)\n",
            "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3->open3d) (3.1.4)\n",
            "Requirement already satisfied: python-dateutil>=2.7 in /usr/local/lib/python3.10/dist-packages (from matplotlib>=3->open3d) (2.8.2)\n",
            "Requirement already satisfied: fastjsonschema>=2.15 in /usr/local/lib/python3.10/dist-packages (from nbformat>=5.7.0->open3d) (2.20.0)\n",
            "Requirement already satisfied: jsonschema>=2.6 in /usr/local/lib/python3.10/dist-packages (from nbformat>=5.7.0->open3d) (4.23.0)\n",
            "Requirement already satisfied: jupyter-core!=5.0.*,>=4.12 in /usr/local/lib/python3.10/dist-packages (from nbformat>=5.7.0->open3d) (5.7.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.0->open3d) (2024.1)\n",
            "Requirement already satisfied: tzdata>=2022.1 in /usr/local/lib/python3.10/dist-packages (from pandas>=1.0->open3d) (2024.1)\n",
            "Requirement already satisfied: scipy>=1.5.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn>=0.21->open3d) (1.13.1)\n",
            "Requirement already satisfied: joblib>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from scikit-learn>=0.21->open3d) (1.4.2)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn>=0.21->open3d) (3.5.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.1.1 in /usr/local/lib/python3.10/dist-packages (from werkzeug>=2.2.3->open3d) (2.1.5)\n",
            "Requirement already satisfied: Jinja2>=3.0 in /usr/local/lib/python3.10/dist-packages (from Flask<3.1,>=1.0.4->dash>=2.6.0->open3d) (3.1.4)\n",
            "Requirement already satisfied: itsdangerous>=2.0 in /usr/local/lib/python3.10/dist-packages (from Flask<3.1,>=1.0.4->dash>=2.6.0->open3d) (2.2.0)\n",
            "Requirement already satisfied: click>=8.0 in /usr/local/lib/python3.10/dist-packages (from Flask<3.1,>=1.0.4->dash>=2.6.0->open3d) (8.1.7)\n",
            "Collecting jedi>=0.16 (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d)\n",
            "  Using cached jedi-0.19.1-py2.py3-none-any.whl.metadata (22 kB)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (4.4.2)\n",
            "Requirement already satisfied: pickleshare in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (0.7.5)\n",
            "Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (3.0.47)\n",
            "Requirement already satisfied: pygments in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (2.16.1)\n",
            "Requirement already satisfied: backcall in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (0.2.0)\n",
            "Requirement already satisfied: matplotlib-inline in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (0.1.7)\n",
            "Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.10/dist-packages (from ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (4.9.0)\n",
            "Requirement already satisfied: attrs>=22.2.0 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=2.6->nbformat>=5.7.0->open3d) (24.2.0)\n",
            "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=2.6->nbformat>=5.7.0->open3d) (2023.12.1)\n",
            "Requirement already satisfied: referencing>=0.28.4 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=2.6->nbformat>=5.7.0->open3d) (0.35.1)\n",
            "Requirement already satisfied: rpds-py>=0.7.1 in /usr/local/lib/python3.10/dist-packages (from jsonschema>=2.6->nbformat>=5.7.0->open3d) (0.20.0)\n",
            "Requirement already satisfied: platformdirs>=2.5 in /usr/local/lib/python3.10/dist-packages (from jupyter-core!=5.0.*,>=4.12->nbformat>=5.7.0->open3d) (4.2.2)\n",
            "Requirement already satisfied: tenacity>=6.2.0 in /usr/local/lib/python3.10/dist-packages (from plotly>=5.0.0->dash>=2.6.0->open3d) (9.0.0)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.7->matplotlib>=3->open3d) (1.16.0)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.10/dist-packages (from importlib-metadata->dash>=2.6.0->open3d) (3.20.1)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->dash>=2.6.0->open3d) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->dash>=2.6.0->open3d) (3.8)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->dash>=2.6.0->open3d) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->dash>=2.6.0->open3d) (2024.8.30)\n",
            "Requirement already satisfied: parso<0.9.0,>=0.8.3 in /usr/local/lib/python3.10/dist-packages (from jedi>=0.16->ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (0.8.4)\n",
            "Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.10/dist-packages (from pexpect>4.3->ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (0.7.0)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.10/dist-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython>=6.1.0->ipywidgets>=8.0.4->open3d) (0.2.13)\n",
            "Downloading open3d-0.18.0-cp310-cp310-manylinux_2_27_x86_64.whl (399.7 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m399.7/399.7 MB\u001b[0m \u001b[31m1.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading dash-2.18.0-py3-none-any.whl (7.5 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.5/7.5 MB\u001b[0m \u001b[31m61.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading dash_core_components-2.0.0-py3-none-any.whl (3.8 kB)\n",
            "Downloading dash_html_components-2.0.0-py3-none-any.whl (4.1 kB)\n",
            "Downloading dash_table-5.0.0-py3-none-any.whl (3.9 kB)\n",
            "Downloading ipywidgets-8.1.5-py3-none-any.whl (139 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m139.8/139.8 kB\u001b[0m \u001b[31m9.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading addict-2.4.0-py3-none-any.whl (3.8 kB)\n",
            "Downloading ConfigArgParse-1.7-py3-none-any.whl (25 kB)\n",
            "Downloading pyquaternion-0.9.9-py3-none-any.whl (14 kB)\n",
            "Downloading comm-0.2.2-py3-none-any.whl (7.2 kB)\n",
            "Downloading widgetsnbextension-4.0.13-py3-none-any.whl (2.3 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.3/2.3 MB\u001b[0m \u001b[31m51.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading retrying-1.3.4-py3-none-any.whl (11 kB)\n",
            "Using cached jedi-0.19.1-py2.py3-none-any.whl (1.6 MB)\n",
            "Installing collected packages: dash-table, dash-html-components, dash-core-components, addict, widgetsnbextension, retrying, pyquaternion, jedi, configargparse, comm, ipywidgets, dash, open3d\n",
            "  Attempting uninstall: widgetsnbextension\n",
            "    Found existing installation: widgetsnbextension 3.6.8\n",
            "    Uninstalling widgetsnbextension-3.6.8:\n",
            "      Successfully uninstalled widgetsnbextension-3.6.8\n",
            "  Attempting uninstall: ipywidgets\n",
            "    Found existing installation: ipywidgets 7.7.1\n",
            "    Uninstalling ipywidgets-7.7.1:\n",
            "      Successfully uninstalled ipywidgets-7.7.1\n",
            "Successfully installed addict-2.4.0 comm-0.2.2 configargparse-1.7 dash-2.18.0 dash-core-components-2.0.0 dash-html-components-2.0.0 dash-table-5.0.0 ipywidgets-8.1.5 jedi-0.19.1 open3d-0.18.0 pyquaternion-0.9.9 retrying-1.3.4 widgetsnbextension-4.0.13\n",
            "Collecting pywavefront\n",
            "  Downloading PyWavefront-1.3.3-py3-none-any.whl.metadata (8.3 kB)\n",
            "Downloading PyWavefront-1.3.3-py3-none-any.whl (28 kB)\n",
            "Installing collected packages: pywavefront\n",
            "Successfully installed pywavefront-1.3.3\n",
            "Collecting trimesh\n",
            "  Downloading trimesh-4.4.9-py3-none-any.whl.metadata (18 kB)\n",
            "Requirement already satisfied: numpy>=1.20 in /usr/local/lib/python3.10/dist-packages (from trimesh) (1.26.4)\n",
            "Downloading trimesh-4.4.9-py3-none-any.whl (700 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m700.1/700.1 kB\u001b[0m \u001b[31m9.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: trimesh\n",
            "Successfully installed trimesh-4.4.9\n"
          ]
        }
      ],
      "source": [
        "!pip3 install open3d\n",
        "!pip3 install pywavefront\n",
        "!pip3 install trimesh\n",
        "!pip3 install timm\n",
        "!pip3 install einops\n",
        "!pip3 install unzip\n",
        "!pip3 install wget\n",
        "!pip3 install vtk"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# 加载最优训练权重 Pre & Vel\n",
        "transolver"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "'wget' �����ڲ����ⲿ���Ҳ���ǿ����еĳ���\n",
            "���������ļ���\n"
          ]
        }
      ],
      "source": [
        "# !wget --load-cookies /tmp/cookies.txt \"https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=FILEID' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\\1\\n/p')&id=FILEID\" -O FILENAME && rm -rf /tmp/cookies.txt\n",
        "# best pre: https://drive.google.com/file/d/1lbvBLnjb2Qu3CCOWX_czxregOMCkp5bv/view?usp=drive_link\n",
        "# bset vel: https://drive.google.com/file/d/1B6vfnDmk2OUYghzRVknvUkOangZ91EmB/view?usp=sharing\n",
        "# https://drive.google.com/uc?id=\n",
        "# Download pre\n",
        "!gdown https://drive.google.com/uc?id=1lbvBLnjb2Qu3CCOWX_czxregOMCkp5bv\n",
        "# Download vel\n",
        "!gdown https://drive.google.com/uc?id=1B6vfnDmk2OUYghzRVknvUkOangZ91EmB"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xgKMibfgYehJ"
      },
      "source": [
        "# 代码运行与数据读取样例\n",
        "（所有训练和推理过程需要在本部分完成）"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RtWKnq1LYlbo"
      },
      "source": [
        "## Dataset1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 读取数据集pre"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IhmnUP3AYk8Z",
        "outputId": "0b7b5818-8a9f-4981-c215-1372b48e877d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Number of points: 29498\n",
            "Number of cells: 26112\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "import open3d\n",
        "import os\n",
        "import torch\n",
        "from torch.utils.data import Dataset\n",
        "\n",
        "\n",
        "# 按照顺序读取标签\n",
        "def read_order(path):\n",
        "    path = os.path.join(path, 'watertight_meshes.txt')\n",
        "    order = []\n",
        "    with open(path,mode='r') as f:\n",
        "        lines = f.readlines()\n",
        "        for line in lines:\n",
        "            # 需要清除换行符\n",
        "            order.append(line.strip())\n",
        "    return order        \n",
        "\n",
        "# 加载特征和标签的路径\n",
        "def load_train(path=r'Training/Dataset_1/'):\n",
        "    n_path = os.path.join(path, 'Feature_File')\n",
        "    order = read_order(path=n_path)\n",
        "    feature_path = []\n",
        "    label_path = []\n",
        "    f_base_path = os.path.join(path, 'Feature_File')\n",
        "    l_base_path = os.path.join(path, 'Label_File')\n",
        "    for o in order:\n",
        "        f_path = f_base_path + f\"/mesh_{o}.ply\"\n",
        "        l_path = l_base_path + f\"/press_{o}.npy\"\n",
        "        feature_path.append(f_path)\n",
        "        label_path.append(l_path)\n",
        "    return feature_path, label_path\n",
        "\n",
        "\n",
        "mean = -37.0900\n",
        "std = 48.0955\n",
        "\n",
        "class MeshPressDataset(Dataset):\n",
        "    def __init__(self, feature_dir, press_dir):\n",
        "        # 特征和标签的文件夹\n",
        "        self.feature_dir = feature_dir\n",
        "        self.press_dir = press_dir\n",
        "        # 判读长度是否一致\n",
        "        assert len(self.feature_dir) == len(self.press_dir)\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.feature_dir)\n",
        "    \n",
        "    # 创造一个迭代器获取数据集\n",
        "    def __getitem__(self, idx):\n",
        "        pcd = open3d.io.read_point_cloud(self.feature_dir[idx])\n",
        "        # 转换为np数据\n",
        "        vertices = np.array(pcd.points)\n",
        "        # 在此处进行裁剪压力\n",
        "        # print(self.label_paths[idx])\n",
        "        press = np.load(self.press_dir[idx])\n",
        "        press = np.concatenate((press[:16], press[112:]), axis=0)\n",
        "        press = (press - mean)/std\n",
        "        # 确保数据长度一致\n",
        "        assert vertices.shape[0] == press.shape[0]\n",
        "\n",
        "        return torch.tensor(vertices, dtype=torch.float32), torch.tensor(press, dtype=torch.float32)\n",
        "\n",
        "\n",
        "\n",
        "def read_test(path=r'Test/Dataset_1/Feature_File'):\n",
        "    \"\"\"需要传入测试集文件夹路径\"\"\"\n",
        "    t_base_path = path\n",
        "    paths = os.listdir(t_base_path)\n",
        "    test_path = []\n",
        "    for p in paths:\n",
        "        if p.split(\".\")[1] == 'ply':\n",
        "            t_path = t_base_path + p\n",
        "            test_path.append(t_path)\n",
        "    return test_path\n",
        "\n",
        "class TestDataset(Dataset):\n",
        "    def __init__(self, test_dir):\n",
        "        self.test_dir = test_dir\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.test_dir)\n",
        "    \n",
        "    def __getitem__(self, idx):\n",
        "        pcd = open3d.io.read_point_cloud(self.test_dir[idx])\n",
        "        # 转换为np数据\n",
        "        vertices = np.array(pcd.points)\n",
        "\n",
        "        return torch.tensor(vertices, dtype=torch.float32), self.test_dir[idx]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 定义损失函数"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# 加载损失函数\n",
        "import torch.nn as nn\n",
        "import torch\n",
        "\n",
        "# 定义损失函数\n",
        "class LpLoss(nn.Module):\n",
        "    def __init__(self, d=2, p=2, size_average=True, reduction='mean'):\n",
        "        super(LpLoss, self).__init__()\n",
        "        assert d > 0 and p > 0\n",
        "        self.d = d\n",
        "        self.p = p\n",
        "        self.reduction = reduction\n",
        "        self.size_average = size_average\n",
        "\n",
        "    def forward(self, x, y):\n",
        "        num_examples = x.size(0)\n",
        "        h = 1.0 / (x.size(1) - 1.0)\n",
        "        all_norms = (h ** (self.d / self.p)) * torch.norm(\n",
        "            x.reshape(num_examples, -1) - y.reshape(num_examples, -1), self.p, dim=1\n",
        "        )\n",
        "        if self.reduction == 'mean':\n",
        "            return all_norms.mean() if self.size_average else all_norms.sum()\n",
        "        elif self.reduction == 'sum':\n",
        "            return all_norms.sum()\n",
        "        else:\n",
        "            return all_norms"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Transolver模型"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import torch.nn as nn\n",
        "from timm.models.layers import trunc_normal_\n",
        "from einops import rearrange, repeat\n",
        "\n",
        "ACTIVATION = {'gelu': nn.GELU, 'tanh': nn.Tanh, 'sigmoid': nn.Sigmoid, 'relu': nn.ReLU, 'leaky_relu': nn.LeakyReLU(0.1),\n",
        "              'softplus': nn.Softplus, 'ELU': nn.ELU, 'silu': nn.SiLU}\n",
        "\n",
        "\n",
        "class Physics_Attention_Irregular_Mesh(nn.Module):\n",
        "    def __init__(self, dim, heads=8, dim_head=64, dropout=0., slice_num=64):\n",
        "        super().__init__()\n",
        "        inner_dim = dim_head * heads\n",
        "        self.dim_head = dim_head\n",
        "        self.heads = heads\n",
        "        self.scale = dim_head ** -0.5\n",
        "        self.softmax = nn.Softmax(dim=-1)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        self.temperature = nn.Parameter(torch.ones([1, heads, 1, 1]) * 0.5)\n",
        "\n",
        "        self.in_project_x = nn.Linear(dim, inner_dim)\n",
        "        self.in_project_fx = nn.Linear(dim, inner_dim)\n",
        "        self.in_project_slice = nn.Linear(dim_head, slice_num)\n",
        "        for l in [self.in_project_slice]:\n",
        "            torch.nn.init.orthogonal_(l.weight)  # use a principled initialization\n",
        "        self.to_q = nn.Linear(dim_head, dim_head, bias=False)\n",
        "        self.to_k = nn.Linear(dim_head, dim_head, bias=False)\n",
        "        self.to_v = nn.Linear(dim_head, dim_head, bias=False)\n",
        "        self.to_out = nn.Sequential(\n",
        "            nn.Linear(inner_dim, dim),\n",
        "            nn.Dropout(dropout)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        # B N C\n",
        "        # print(x.shape)\n",
        "        _, B, N, C = x.shape\n",
        "        \n",
        "        ### (1) Slice\n",
        "        fx_mid = self.in_project_fx(x).reshape(B, N, self.heads, self.dim_head) \\\n",
        "            .permute(0, 2, 1, 3).contiguous()  # B H N C\n",
        "        x_mid = self.in_project_x(x).reshape(B, N, self.heads, self.dim_head) \\\n",
        "            .permute(0, 2, 1, 3).contiguous()  # B H N C\n",
        "        slice_weights = self.softmax(self.in_project_slice(x_mid) / self.temperature)  # B H N G\n",
        "        slice_norm = slice_weights.sum(2)  # B H G\n",
        "        slice_token = torch.einsum(\"bhnc,bhng->bhgc\", fx_mid, slice_weights)\n",
        "        slice_token = slice_token / ((slice_norm + 1e-5)[:, :, :, None].repeat(1, 1, 1, self.dim_head))\n",
        "\n",
        "        ### (2) Attention among slice tokens\n",
        "        q_slice_token = self.to_q(slice_token)\n",
        "        k_slice_token = self.to_k(slice_token)\n",
        "        v_slice_token = self.to_v(slice_token)\n",
        "        dots = torch.matmul(q_slice_token, k_slice_token.transpose(-1, -2)) * self.scale\n",
        "        attn = self.softmax(dots)\n",
        "        attn = self.dropout(attn)\n",
        "        out_slice_token = torch.matmul(attn, v_slice_token)  # B H G D\n",
        "\n",
        "        ### (3) Deslice\n",
        "        out_x = torch.einsum(\"bhgc,bhng->bhnc\", out_slice_token, slice_weights)\n",
        "        out_x = rearrange(out_x, 'b h n d -> b n (h d)')\n",
        "        return self.to_out(out_x)\n",
        "\n",
        "\n",
        "class MLP(nn.Module):\n",
        "    def __init__(self, n_input, n_hidden, n_output, n_layers=1, act='gelu', res=True):\n",
        "        super(MLP, self).__init__()\n",
        "\n",
        "        if act in ACTIVATION.keys():\n",
        "            act = ACTIVATION[act]\n",
        "        else:\n",
        "            raise NotImplementedError\n",
        "        self.n_input = n_input\n",
        "        self.n_hidden = n_hidden\n",
        "        self.n_output = n_output\n",
        "        self.n_layers = n_layers\n",
        "        self.res = res\n",
        "        self.linear_pre = nn.Sequential(nn.Linear(n_input, n_hidden), act())\n",
        "        self.linear_post = nn.Linear(n_hidden, n_output)\n",
        "        self.linears = nn.ModuleList([nn.Sequential(nn.Linear(n_hidden, n_hidden), act()) for _ in range(n_layers)])\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.linear_pre(x)\n",
        "        for i in range(self.n_layers):\n",
        "            if self.res:\n",
        "                x = self.linears[i](x) + x\n",
        "            else:\n",
        "                x = self.linears[i](x)\n",
        "        x = self.linear_post(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "class Transolver_block(nn.Module):\n",
        "    \"\"\"Transformer encoder block.\"\"\"\n",
        "\n",
        "    def __init__(\n",
        "            self,\n",
        "            num_heads: int,\n",
        "            hidden_dim: int,\n",
        "            dropout: float,\n",
        "            act='gelu',\n",
        "            mlp_ratio=4,\n",
        "            last_layer=False,\n",
        "            out_dim=1,\n",
        "            slice_num=32,\n",
        "    ):\n",
        "        super().__init__()\n",
        "        self.last_layer = last_layer\n",
        "        self.ln_1 = nn.LayerNorm(hidden_dim)\n",
        "        self.Attn = Physics_Attention_Irregular_Mesh(hidden_dim, heads=num_heads, dim_head=hidden_dim // num_heads,\n",
        "                                                     dropout=dropout, slice_num=slice_num)\n",
        "        self.ln_2 = nn.LayerNorm(hidden_dim)\n",
        "        self.mlp = MLP(hidden_dim, hidden_dim * mlp_ratio, hidden_dim, n_layers=0, res=False, act=act)\n",
        "        if self.last_layer:\n",
        "            self.ln_3 = nn.LayerNorm(hidden_dim)\n",
        "            self.mlp2 = nn.Linear(hidden_dim, out_dim)\n",
        "\n",
        "    def forward(self, fx):\n",
        "        fx = self.Attn(self.ln_1(fx)) + fx\n",
        "        fx = self.mlp(self.ln_2(fx)) + fx\n",
        "        if self.last_layer:\n",
        "            return self.mlp2(self.ln_3(fx))\n",
        "        else:\n",
        "            return fx\n",
        "\n",
        "\n",
        "class Model(nn.Module):\n",
        "    def __init__(self,\n",
        "                 space_dim=1,\n",
        "                 n_layers=5,\n",
        "                 n_hidden=256,\n",
        "                 dropout=0,\n",
        "                 n_head=8,\n",
        "                 act='gelu',\n",
        "                 mlp_ratio=1,\n",
        "                 fun_dim=1,\n",
        "                 out_dim=1,\n",
        "                 slice_num=32,\n",
        "                 ref=8,\n",
        "                 unified_pos=False\n",
        "                 ):\n",
        "        super(Model, self).__init__()\n",
        "        self.__name__ = 'Transolver'\n",
        "        self.ref = ref\n",
        "        self.unified_pos = unified_pos\n",
        "        if self.unified_pos:\n",
        "            self.preprocess = MLP(fun_dim + self.ref * self.ref * self.ref, n_hidden * 2, n_hidden, n_layers=0,\n",
        "                                  res=False, act=act)\n",
        "        else:\n",
        "            self.preprocess = MLP(fun_dim + space_dim, n_hidden * 2, n_hidden, n_layers=0, res=False, act=act)\n",
        "\n",
        "        self.n_hidden = n_hidden\n",
        "        self.space_dim = space_dim\n",
        "\n",
        "        self.blocks = nn.ModuleList([Transolver_block(num_heads=n_head, hidden_dim=n_hidden,\n",
        "                                                      dropout=dropout,\n",
        "                                                      act=act,\n",
        "                                                      mlp_ratio=mlp_ratio,\n",
        "                                                      out_dim=out_dim,\n",
        "                                                      slice_num=slice_num,\n",
        "                                                      last_layer=(_ == n_layers - 1))\n",
        "                                     for _ in range(n_layers)])\n",
        "        self.initialize_weights()\n",
        "        self.placeholder = nn.Parameter((1 / (n_hidden)) * torch.rand(n_hidden, dtype=torch.float))\n",
        "\n",
        "    def initialize_weights(self):\n",
        "        self.apply(self._init_weights)\n",
        "\n",
        "    def _init_weights(self, m):\n",
        "        if isinstance(m, nn.Linear):\n",
        "            trunc_normal_(m.weight, std=0.02)\n",
        "            if isinstance(m, nn.Linear) and m.bias is not None:\n",
        "                nn.init.constant_(m.bias, 0)\n",
        "        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d)):\n",
        "            nn.init.constant_(m.bias, 0)\n",
        "            nn.init.constant_(m.weight, 1.0)\n",
        "\n",
        "    def get_grid(self, my_pos):\n",
        "        # my_pos 1 N 3\n",
        "        batchsize = my_pos.shape[0]\n",
        "\n",
        "        gridx = torch.tensor(np.linspace(-1.5, 1.5, self.ref), dtype=torch.float)\n",
        "        gridx = gridx.reshape(1, self.ref, 1, 1, 1).repeat([batchsize, 1, self.ref, self.ref, 1])\n",
        "        gridy = torch.tensor(np.linspace(0, 2, self.ref), dtype=torch.float)\n",
        "        gridy = gridy.reshape(1, 1, self.ref, 1, 1).repeat([batchsize, self.ref, 1, self.ref, 1])\n",
        "        gridz = torch.tensor(np.linspace(-4, 4, self.ref), dtype=torch.float)\n",
        "        gridz = gridz.reshape(1, 1, 1, self.ref, 1).repeat([batchsize, self.ref, self.ref, 1, 1])\n",
        "        grid_ref = torch.cat((gridx, gridy, gridz), dim=-1).cuda().reshape(batchsize, self.ref ** 3, 3)  # B 4 4 4 3\n",
        "\n",
        "        pos = torch.sqrt(\n",
        "            torch.sum((my_pos[:, :, None, :] - grid_ref[:, None, :, :]) ** 2,\n",
        "                      dim=-1)). \\\n",
        "            reshape(batchsize, my_pos.shape[1], self.ref * self.ref * self.ref).contiguous()\n",
        "        return pos\n",
        "\n",
        "    def forward(self, data):\n",
        "        cfd_data = data\n",
        "        x, fx = cfd_data, None\n",
        "        x = x[None, :, :]\n",
        "        # 没有unified_pos\n",
        "        if self.unified_pos:\n",
        "            new_pos = self.get_grid(cfd_data.pos[None, :, :])\n",
        "            x = torch.cat((x, new_pos), dim=-1)\n",
        "\n",
        "        if fx is not None:\n",
        "            fx = torch.cat((x, fx), -1)\n",
        "            fx = self.preprocess(fx)\n",
        "        else:\n",
        "            fx = self.preprocess(x)\n",
        "            fx = fx + self.placeholder[None, None, :]\n",
        "\n",
        "        for block in self.blocks:\n",
        "            fx = block(fx)\n",
        "\n",
        "        return fx[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 定义超参数 Pressure"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# 获取模型超参数\n",
        "import argparse\n",
        "\n",
        "def get_args():\n",
        "    parser = argparse.ArgumentParser(description=\"use config param for Transovler 9.15\")\n",
        "\n",
        "    # --------------------初始化各个参数------------------------------\n",
        "    # training Setting\n",
        "    parser.add_argument('--epochs', type=int, default=300, help=\"set the train epochs\")\n",
        "    parser.add_argument(\"--lr\", type=float, default=0.001, help=\"maybe 0.001 is large\")\n",
        "    parser.add_argument(\"--seed\", type=int, default=1225, help=\"Birthday\")\n",
        "    parser.add_argument(\"--log_dir\", default=\"trainL40\", help=\"setting the save file name\")\n",
        "    parser.add_argument(\"--batch\", type=int, default=2, help=\"for training 2 is best\")\n",
        "    parser.add_argument(\"--train_ratio\", type=float, default=0.98, help=\"500*0.98 -> 490:10\")\n",
        "    parser.add_argument(\"--val_batch\", type=int, default=4, help=\"default=1\")\n",
        "    parser.add_argument(\"--max_lr\", type=float, default=0.001, help=\"use OneCycle method\")\n",
        "    parser.add_argument('--train_path', default='Training/Dataset_1/Feature_File')\n",
        "    # model Setting\n",
        "    parser.add_argument(\"--n_layers\", type=int, default=40, help=\"maybe 7 is best\")\n",
        "    parser.add_argument(\"--n_hidden\", type=int, default=256, help=\"\")\n",
        "    parser.add_argument(\"--space_dim\", type=int, default=3, help=\"input data's dim, in trackA is 3\")\n",
        "    parser.add_argument(\"--mlp_ratio\", type=int, default=2, help='2 is best? maybe')\n",
        "    parser.add_argument(\"--fun_dim\", type=int, default=0, help=\"set 0 !!\")\n",
        "    parser.add_argument(\"--n_head\", type=int, default=8, help='atten head')\n",
        "    parser.add_argument(\"--slice_num\", type=int, default=512, help='physic atten slice')\n",
        "    parser.add_argument(\"--out_dim\", type=int, default=1, help='output')\n",
        "    parser.add_argument(\"--frezze\", action='store_true', help='frezze some params')\n",
        "    parser.add_argument('--fre_layers', type=list, default=None, help='choose fre some layer')\n",
        "    parser.add_argument('--fre_epoch', type=int, default=220, help='star frezee')\n",
        "    parser.add_argument(\"--act\", default='gelu', help=\"activate function\", choices=['gelu', 'tanh', 'sigmoid', 'relu', 'leaky_relu', 'softplus', 'ELU', 'silu'])\n",
        "\n",
        "    parser.add_argument('--result', default='best', choices=['best', 'last'], help=\"use best or last weights\")\n",
        "    parser.add_argument('--output_file', default=\"submission\", help='need to submit')\n",
        "    parser.add_argument('--is_zip', default=1, choices=[0, 1], help='zip or not zip')\n",
        "    parser.add_argument('--test_path', default=r'Test/Dataset_1/Feature_File/', help='TestSet Path')\n",
        "\n",
        "    # opt = parser.parse_args()\n",
        "    opt = parser.parse_known_args()[0]\n",
        "\n",
        "    return opt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 训练模型，并保存超参数到runs_trackA/train100目录下"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# -----------------------训练模型------------------------------\n",
        "import os\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch\n",
        "from torch.utils.data import random_split,DataLoader\n",
        "from time import time\n",
        "from tqdm import tqdm\n",
        "import numpy as np\n",
        "import random\n",
        "import csv\n",
        "\n",
        "\n",
        "mean = -37.0900\n",
        "std = 48.0955\n",
        "\n",
        "def set_seed(seed=1225):\n",
        "    # 固定随机种子\n",
        "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
        "    random.seed(seed)\n",
        "    np.random.seed(seed)\n",
        "    torch.manual_seed(seed)\n",
        "    torch.cuda.manual_seed(seed)\n",
        "    torch.backends.cudnn.deterministic = True\n",
        "    torch.backends.cudnn.benchmark = False\n",
        "    msg =f\"\"\"\n",
        "    =============================================\n",
        "        seed = {seed}\n",
        "        torch.backends.cudnn.deterministic = True\n",
        "        torch.backends.cudnn.benchmark = False\n",
        "    =============================================\n",
        "    \"\"\"\n",
        "    print(msg)\n",
        "    \n",
        "def judge_file(folder_path='runs_trackA'):\n",
        "    folder_path = folder_path\n",
        "\n",
        "    if os.path.exists(folder_path) and os.path.isdir(folder_path):\n",
        "        print(\"文件夹存在\")\n",
        "    else:\n",
        "        os.system(f\"mkdir -p {folder_path}\")\n",
        "\n",
        "def train(opt):\n",
        "\n",
        "    log_p = opt.log_dir\n",
        "    # 设置随机种子\n",
        "    set_seed(opt.seed)\n",
        "    train_split = opt.train_ratio\n",
        "    train_batch = opt.batch\n",
        "    val_batch = opt.val_batch\n",
        "    num_epoch = opt.epochs\n",
        "    n_layers = opt.n_layers\n",
        "    n_hidden = opt.n_hidden\n",
        "    space_dim = opt.space_dim\n",
        "    mlp_ratio = opt.mlp_ratio\n",
        "    fun_dim = opt.fun_dim\n",
        "    lr = opt.lr\n",
        "    max_lr = opt.max_lr\n",
        "    act = opt.act\n",
        "    slice_num = opt.slice_num\n",
        "    n_head = opt.n_head\n",
        "    out_dim = opt.out_dim\n",
        "    \n",
        "    # 创建保存权重的文件夹\n",
        "    save_file = 'runs_trackA'\n",
        "    judge_file(save_file)\n",
        "    log_file = os.path.join(save_file, log_p)\n",
        "    judge_file(log_file)\n",
        "\n",
        "    # 保存训练过程到csv文件中\n",
        "    f = open(f'{log_file}/train_process.csv', mode='w')\n",
        "    fc = csv.writer(f)\n",
        "    fc.writerow([\"train_loss\", \"val_loss\", 'lr'])\n",
        "\n",
        "    with open(f'{log_file}/args.txt', mode='w') as fargs:\n",
        "        fargs.write(str(opt))\n",
        "    \n",
        "    print(f\"saving args to {log_file}\")\n",
        "    \n",
        "    # 1. -加载数据集 百度给的数据集\n",
        "    feature_paths, label_paths = load_train()\n",
        "    dataset = MeshPressDataset(feature_paths, label_paths)\n",
        "    \n",
        "\n",
        "    # 检查gpu是否可用\n",
        "    if torch.cuda.is_available():\n",
        "        device = torch.device(\"cuda\")\n",
        "        print(\"GPU is available\")\n",
        "    else:\n",
        "        device = torch.device(\"cpu\")\n",
        "        print(\"GPU is not available, using CPU\")\n",
        "\n",
        "    \n",
        "    # 2. -划分数据集为训练集和验证集 0.8 : 0.2\n",
        "    # 可以试一下 0.9:0.1 可能效果会好很多\n",
        "    train_size = int(train_split * len(dataset))\n",
        "    val_size = len(dataset) - train_size\n",
        "    train_dataset, val_dataset = random_split(dataset, [train_size, val_size])\n",
        "\n",
        "\n",
        "\n",
        "    # 3. -创建数据集加载器\n",
        "    train_loader = DataLoader(train_dataset, batch_size=train_batch, shuffle=True)\n",
        "    val_loader = DataLoader(val_dataset, batch_size=val_batch, shuffle=False)\n",
        "\n",
        "\n",
        "    # 4. -实例化模型\n",
        "    model = Model(\n",
        "        n_layers=n_layers, \n",
        "        n_hidden=n_hidden, \n",
        "        space_dim=space_dim, \n",
        "        mlp_ratio=mlp_ratio, \n",
        "        fun_dim=fun_dim,\n",
        "        act=act,\n",
        "        n_head=n_head,\n",
        "        slice_num=slice_num,\n",
        "        out_dim = out_dim\n",
        "        )\n",
        "\n",
        "        \n",
        "    \n",
        "    model.to(device)\n",
        "\n",
        "    params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
        "    print(f\"---模型的参数为：{params} ---\")\n",
        "    \n",
        "\n",
        "    # 5. -定义损失函数\n",
        "    criterion = LpLoss(size_average=True)\n",
        "\n",
        "    # 6. -定义优化器 修改不同的参数，就不能加载预训练权重，会重新开始训练，最好的方法就是修改存储weights的文件夹的名称 学习率不能太大，太难了乱跑 学习率需要小一些 下降的速度也需要慢一点，太快了容易提前收敛了\n",
        "    #  之后需要修改lr lr_decay,获取弄的比较好的效果\n",
        "    # optimizer = optim.Adam(model.parameters(), lr=0.0001)\n",
        "    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.00005, betas=(0.9, 0.999))\n",
        "    \n",
        "    # 7. -开始训练\n",
        "    best_val_loss = 1000\n",
        "    best_train_loss = 1000\n",
        "    # 把学习率下降调小\n",
        "\n",
        "    scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
        "        optimizer,\n",
        "        max_lr=max_lr,\n",
        "        epochs=num_epoch,\n",
        "        div_factor=1e4, pct_start=0.25, final_div_factor=1e4,\n",
        "        steps_per_epoch=len(train_loader)\n",
        "    )\n",
        "    \n",
        "    # scheduler = optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.000001, max_lr=0.0008, step_size_up=800, mode='triangular')\n",
        "    grad_clip = 1000.0\n",
        "    for epoch in range(num_epoch):\n",
        "        t1 = time()\n",
        "        \n",
        "        \n",
        "        model.train()\n",
        "        train_loss = 0\n",
        "        for X_batch, y_batch in tqdm(train_loader):\n",
        "            X_batch, y_batch = X_batch.to(device), y_batch.to(device)\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "            outputs = model(X_batch)\n",
        "\n",
        "            outputs = outputs * std + mean\n",
        "            y_batch = y_batch * std + mean\n",
        "            \n",
        "            loss = criterion(outputs,y_batch)\n",
        "            \n",
        "            loss.backward()\n",
        "            nn.utils.clip_grad_norm_(model.parameters(), grad_clip)\n",
        "            optimizer.step()\n",
        "            scheduler.step()\n",
        "            \n",
        "            \n",
        "            train_loss += loss.item()\n",
        "            # break\n",
        "        train_loss /= len(train_loader)\n",
        "\n",
        "        model.eval()\n",
        "        val_loss = 0\n",
        "        with torch.no_grad():\n",
        "            for X_batch, y_batch in val_loader:\n",
        "                X_batch, y_batch = X_batch.to(device), y_batch.to(device)\n",
        "                outputs = model(X_batch)\n",
        "                outputs = outputs * std + mean\n",
        "                y_batch = y_batch * std + mean\n",
        "                loss = criterion(outputs, y_batch)\n",
        "                val_loss += loss.item()\n",
        "        val_loss /= len(val_loader)\n",
        "\n",
        "        \n",
        "\n",
        "        # 8. -保存模型最后一轮权重\n",
        "        last_weights_path = f'{log_file}/last_weights.pt'\n",
        "        torch.save(model.state_dict(),last_weights_path)\n",
        "\n",
        "        # 9. -保存最优权重\n",
        "        if val_loss <= best_val_loss and train_loss <= best_train_loss:\n",
        "            best_val_loss = val_loss\n",
        "            best_train_loss = train_loss\n",
        "            best_weight_path = f'{log_file}/best_weights.pt'\n",
        "            torch.save(model.state_dict(), best_weight_path)\n",
        "        \n",
        "        current_lr = optimizer.param_groups[0]['lr']\n",
        "\n",
        "        \n",
        "        if epoch+1 == opt.fre_epoch:\n",
        "            if opt.frezze:\n",
        "                # if opt.fre_epoch <= opt.epochs:\n",
        "                \n",
        "                for param in model.parameters():\n",
        "                    param.requires_grad = True\n",
        "                \n",
        "                for fre_layer in opt.fre_layers:\n",
        "                    # print(fre_layer)\n",
        "                    # 冻结选择的层数\n",
        "                    for param in model.blocks[int(fre_layer)].parameters():\n",
        "                        param.requires_grad = False\n",
        "                optimizer = torch.optim.AdamW(model.parameters(), lr=current_lr, weight_decay=0.00005, betas=(0.9, 0.999))\n",
        "                scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
        "                    optimizer,\n",
        "                    max_lr=current_lr,\n",
        "                    epochs=num_epoch-opt.fre_epoch,\n",
        "                    div_factor=1e4, pct_start=0.001, final_div_factor=1e4,\n",
        "                    steps_per_epoch=len(train_loader)\n",
        "                )\n",
        "        \n",
        "        t2 =time()\n",
        "        t = t2 - t1\n",
        "        print(f'Epoch {epoch+1}/{num_epoch}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Current_lr: {current_lr}, Spend Time: {round(t,2)}s')\n",
        "        fc.writerow([train_loss, val_loss, current_lr])\n",
        "\n",
        "    f.close()\n",
        "        "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 运行测试集代码"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# -----------------------------测试结果-------------------------------\n",
        "\n",
        "import numpy as np\n",
        "import os\n",
        "import torch\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "\n",
        "def test(opt, flag='test'):\n",
        "    log = opt.log_dir\n",
        "    mode = opt.result\n",
        "    output_file = opt.output_file\n",
        "    # 检查gpu是否可用\n",
        "    if torch.cuda.is_available():\n",
        "        device = torch.device(\"cuda\")\n",
        "        print(\"GPU is available\")\n",
        "    else:\n",
        "        device = torch.device(\"cpu\")\n",
        "        print(\"GPU is not available, using CPU\")\n",
        "\n",
        "    # 加载模型，进行测试\n",
        "    n_layers = opt.n_layers\n",
        "    n_hidden = opt.n_hidden\n",
        "    space_dim = opt.space_dim\n",
        "    mlp_ratio = opt.mlp_ratio\n",
        "    fun_dim = opt.fun_dim\n",
        "    act = opt.act\n",
        "    slice_num = opt.slice_num\n",
        "    n_head = opt.n_head\n",
        "    model = Model(\n",
        "        n_layers=n_layers, \n",
        "        n_hidden=n_hidden, \n",
        "        space_dim=space_dim, \n",
        "        mlp_ratio=mlp_ratio, \n",
        "        fun_dim=fun_dim,\n",
        "        act=act,\n",
        "        n_head=n_head,\n",
        "        slice_num=slice_num,\n",
        "        )\n",
        "    \n",
        "    \n",
        "    os.system(f'mkdir -p {output_file}')\n",
        "\n",
        "    # 加载最优的权重模型\n",
        "\n",
        "    if flag == 'train':\n",
        "        best_weights_path = f'runs_trackA/{log}/best_weights.pt'\n",
        "    else:\n",
        "        best_weights_path = r'transolver_pre.pt'  # 直接在当前目录下读取最佳权重\n",
        "    model.load_state_dict(torch.load(best_weights_path))\n",
        "    \n",
        "      \n",
        "    # 加载到GPU上\n",
        "    model.to(device)\n",
        "\n",
        "    \n",
        "    model.eval()\n",
        "    # 获取测试集的路径\n",
        "    test_path = read_test(path=opt.test_path)\n",
        "    # 加载测试集数据\n",
        "    dataset = TestDataset(test_dir=test_path)\n",
        "    test_loader = DataLoader(dataset, batch_size=1, shuffle=False)\n",
        "\n",
        "    # 进行测试\n",
        "    for X_batch, p in test_loader:\n",
        "\n",
        "        # 将数据部署到GPU上\n",
        "        X_batch= X_batch.to(device)        \n",
        "        outputs = model(X_batch)\n",
        "        outputs = outputs * std + mean\n",
        "\n",
        "\n",
        "        # 将输出到cpu上\n",
        "        outputs = outputs.to('cpu')\n",
        "\n",
        "        # print(type(outputs))\n",
        "        # 数据进行展平，并转换为numpy的格式\n",
        "        outputs = outputs.detach().numpy().flatten()\n",
        "        print(np.max(outputs))\n",
        "\n",
        "        # 处理保存数据的路径 \n",
        "        number = str(p).split(\".\")[0][-3:]\n",
        "        path = f'{output_file}/press_{number}.npy'\n",
        "\n",
        "        # 保存数据\n",
        "        np.save(path,outputs)\n",
        "\n",
        "    print(\"-----完成数据的保存啦～=～-------\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 测试最优权重，从google drive下载"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "if __name__ == \"__main__\":\n",
        "\n",
        "    opt = get_args()\n",
        "    test(opt)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 训练主函数，并测试验证"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "if __name__ == \"__main__\":\n",
        "\n",
        "    opt = get_args()\n",
        "    train(opt=opt)\n",
        "    test(opt=opt, flag='train')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 速度"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 读取数据集vel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "\n",
        "import numpy as np\n",
        "import open3d\n",
        "import os\n",
        "import vtk\n",
        "from vtk.util.numpy_support import vtk_to_numpy\n",
        "from pathlib import Path\n",
        "from torch.utils.data import Dataset ,DataLoader, random_split\n",
        "from time import time\n",
        "import torch\n",
        "\n",
        "def read_dir(test=False):\n",
        "    vel_path = []\n",
        "    mesh_path = []\n",
        "    if test:\n",
        "        paths = os.listdir(\"Test/Dataset_1/Feature_File/\")\n",
        "        for p in paths:\n",
        "            p = os.path.join(\"Test/Dataset_1/Feature_File/\", p)\n",
        "            p = Path(p)\n",
        "            if p.suffix == '.ply':\n",
        "                mesh_path.append(p)\n",
        "            elif p.suffix == '.vtk':\n",
        "                vel_path.append(p)\n",
        "        return vel_path, mesh_path\n",
        "    else:\n",
        "        paths = os.listdir(\"Training/Dataset_1/Feature_File/\")\n",
        "        # paths = Path(paths)\n",
        "        for p in paths:\n",
        "            p = os.path.join(\"Training/Dataset_1/Feature_File/\", p)\n",
        "            p = Path(p)\n",
        "            if p.suffix == '.ply':\n",
        "                mesh_path.append(p)\n",
        "            elif p.suffix == '.vtk':\n",
        "                vel_path.append(p)\n",
        "\n",
        "\n",
        "        return vel_path, mesh_path\n",
        "\n",
        "\n",
        "\n",
        "# 读取obj数据\n",
        "def read_obj(file_path):\n",
        "    reader = vtk.vtkOBJReader()\n",
        "    reader.SetFileName(file_path)\n",
        "    reader.Update()\n",
        "    polydata = reader.GetOutput()\n",
        "    return reader, polydata\n",
        "\n",
        "\n",
        "# 读取vtk数据\n",
        "def read_vtk(file_path):\n",
        "    reader = vtk.vtkUnstructuredGridReader()\n",
        "    reader.SetFileName(file_path)\n",
        "    reader.Update()\n",
        "    polydata = reader.GetOutput()\n",
        "    return reader, polydata\n",
        "\n",
        "def load_sdf_queries():\n",
        "    tx = np.linspace(0, 1, 64)\n",
        "    ty = np.linspace(0, 1, 64)\n",
        "    tz = np.linspace(0, 1, 64)\n",
        "    sdf_q = np.stack(np.meshgrid(tx, ty, tz, indexing=\"ij\"), axis=-1).astype(np.float32)\n",
        "    sdf_q = np.transpose(sdf_q, (3,0,1,2))\n",
        "    return sdf_q \n",
        "\n",
        "def load_sdf():\n",
        "    sdf = np.ones([64,64,64]).astype(np.float32)\n",
        "    return sdf \n",
        "\n",
        "\n",
        "def normals(polydata):\n",
        "    normals_filter = vtk.vtkPolyDataNormals()\n",
        "    normals_filter.SetInputData(polydata)\n",
        "    normals_filter.ComputeCellNormalsOn()\n",
        "    normals_filter.ConsistencyOn()\n",
        "    normals_filter.FlipNormalsOn()\n",
        "    normals_filter.AutoOrientNormalsOn()\n",
        "    normals_filter.Update()\n",
        "    numpy_cell_normals = vtk_to_numpy(normals_filter.GetOutput().GetCellData().GetNormals()).astype(np.float32)\n",
        "    return numpy_cell_normals\n",
        "\n",
        "def areas(polydata):\n",
        "    cell_size_filter = vtk.vtkCellSizeFilter()\n",
        "    cell_size_filter.SetInputData(polydata)\n",
        "    cell_size_filter.ComputeAreaOn()\n",
        "    cell_size_filter.Update()\n",
        "    numpy_cell_areas = vtk_to_numpy(cell_size_filter.GetOutput().GetCellData().GetArray(\"Area\")).astype(np.float32)\n",
        "    return numpy_cell_areas\n",
        "\n",
        "def centoirds(polydata):\n",
        "    cell_centers = vtk.vtkCellCenters()\n",
        "    cell_centers.SetInputData(polydata)\n",
        "    cell_centers.Update()\n",
        "    numpy_cell_centers = vtk_to_numpy(cell_centers.GetOutput().GetPoints().GetData()).astype(np.float32)\n",
        "    return numpy_cell_centers\n",
        "\n",
        "def nodes(polydata):\n",
        "    # points = vtk_to_numpy(polydata.GetPoints().GetData()).astype(np.float32)\n",
        "    points = vtk_to_numpy(polydata.GetPoints().GetData()).astype(np.float32)\n",
        "    return points\n",
        "\n",
        "def velocity(polydata):\n",
        "    vel = vtk_to_numpy(polydata.GetPointData().GetArray(\"point_vectors\")).astype(np.float32)\n",
        "    return vel\n",
        "\n",
        "\n",
        "# 根据后缀组成不同的数据\n",
        "def read(file_path):\n",
        "    # 读取obj数据集\n",
        "    if file_path.suffix == \".obj\":\n",
        "        _, polydata = read_obj(file_path)\n",
        "        data_dict = {\n",
        "            \"centroids\":    centoirds(polydata),\n",
        "            \"vertices\":     nodes(polydata),\n",
        "            \"areas\":        areas(polydata),\n",
        "            \"normal\":       normals(polydata),\n",
        "            # \"sdf\":          load_sdf(),  去看百度的\n",
        "            # \"sdf_query_points\":          load_sdf_queries(),\n",
        "        }\n",
        "    # 读取速度的数据集\n",
        "    elif file_path.suffix == \".vtk\":\n",
        "        _, polydata = read_vtk(file_path)\n",
        "        data_dict = {\n",
        "            \"vertices\":     nodes(polydata),\n",
        "            \"velocity\":     velocity(polydata),\n",
        "            # \"sdf\":          load_sdf(),\n",
        "            # \"sdf_query_points\":          load_sdf_queries(),\n",
        "        }\n",
        "    else:\n",
        "        raise NotImplemented\n",
        "    \n",
        "    return data_dict\n",
        "\n",
        "def read_test(file_path):\n",
        "    # 读取obj数据集\n",
        "    if file_path.suffix == \".obj\":\n",
        "        _, polydata = read_obj(file_path)\n",
        "        data_dict = {\n",
        "            \"centroids\":    centoirds(polydata),\n",
        "            \"vertices\":     nodes(polydata),\n",
        "            \"areas\":        areas(polydata),\n",
        "            \"normal\":       normals(polydata),\n",
        "            # \"sdf\":          load_sdf(),  去看百度的\n",
        "            # \"sdf_query_points\":          load_sdf_queries(),\n",
        "        }\n",
        "    # 读取速度的数据集\n",
        "    elif file_path.suffix == \".vtk\":\n",
        "        _, polydata = read_vtk(file_path)\n",
        "        data_dict = {\n",
        "            \"vertices\":     nodes(polydata),\n",
        "            # \"velocity\":     velocity(polydata), # 测试集没有速度，需要自己预测\n",
        "            # \"sdf\":          load_sdf(),\n",
        "            # \"sdf_query_points\":          load_sdf_queries(),\n",
        "        }\n",
        "    else:\n",
        "        raise NotImplemented\n",
        "    \n",
        "    return data_dict\n",
        "\n",
        "\n",
        "class VelDataset(Dataset):\n",
        "    def __init__(self, feature_dir):\n",
        "        # 传入一个列表\n",
        "        self.feature_dir = feature_dir\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.feature_dir)\n",
        "    \n",
        "    def __getitem__(self, idx):\n",
        "        p = self.feature_dir[idx]\n",
        "        data = read(p)\n",
        "        return torch.tensor(data['vertices'], dtype=torch.float32), torch.tensor(data['velocity'], dtype=torch.float32)\n",
        "        \n",
        "\n",
        "# 读取速度的测试集\n",
        "class VelTestDataset(Dataset):\n",
        "    def __init__(self, feature_dir):\n",
        "        # 传入一个列表\n",
        "        self.feature_dir = feature_dir\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.feature_dir)\n",
        "    \n",
        "    def __getitem__(self, idx):\n",
        "        p = self.feature_dir[idx]\n",
        "        data = read_test(p)\n",
        "        return torch.tensor(data['vertices'], dtype=torch.float32), str(p)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 定义速度的超参数"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import argparse\n",
        "def get_args():\n",
        "    parser = argparse.ArgumentParser(description=\"use config param for Transovler 7.11\")\n",
        "\n",
        "    # --------------------初始化各个参数------------------------------\n",
        "    # training Setting\n",
        "    parser.add_argument('--epochs', type=int, default=200, help=\"set the train epochs\")\n",
        "    parser.add_argument(\"--lr\", type=float, default=0.001, help=\"maybe 0.001 is large\")\n",
        "    parser.add_argument(\"--seed\", type=int, default=2024, help=\"Birthday\")\n",
        "    parser.add_argument(\"--log_dir\", default=\"trainL16\", help=\"setting the save file name\")\n",
        "    parser.add_argument(\"--batch\", type=int, default=1, help=\"for training\")\n",
        "    parser.add_argument(\"--train_ratio\", type=float, default=0.98, help=\"500*0.98 -> 490:10\")\n",
        "    parser.add_argument(\"--val_batch\", type=int, default=4, help=\"default=1\")\n",
        "    parser.add_argument(\"--max_lr\", type=float, default=0.001, help=\"use OneCycle method\")\n",
        "    # model Setting\n",
        "    parser.add_argument(\"--n_layers\", type=int, default=16, help=\"maybe 8 is best\")\n",
        "    parser.add_argument(\"--n_hidden\", type=int, default=256, help=\"\")\n",
        "    parser.add_argument(\"--space_dim\", type=int, default=3, help=\"input data's dim, in trackA is 3\")\n",
        "    parser.add_argument(\"--mlp_ratio\", type=int, default=2, help='2 is best? maybe')\n",
        "    parser.add_argument(\"--fun_dim\", type=int, default=0, help=\"set 0 !!\")\n",
        "    parser.add_argument(\"--n_head\", type=int, default=8, help='atten head')\n",
        "    parser.add_argument(\"--slice_num\", type=int, default=128, help='physic atten slice')\n",
        "    parser.add_argument('--out_dim', type=int, default=3, help='output dimension')\n",
        "    \n",
        "    parser.add_argument(\"--frezze\", action='store_true', help='frezze some params')\n",
        "    parser.add_argument('--fre_layers', type=list, default=None, help='choose fre some layer')\n",
        "    parser.add_argument('--fre_epoch', type=int, default=300, help='star frezee')\n",
        "    parser.add_argument(\"--act\", default='gelu', help=\"activate function\", choices=['gelu', 'tanh', 'sigmoid', 'relu', 'leaky_relu', 'softplus', 'ELU', 'silu'])\n",
        "\n",
        "    parser.add_argument('--result', default='best', choices=['best', 'last'], help=\"use best or last weights\")\n",
        "    parser.add_argument('--output_file', default=\"submission\", help='need to submit')\n",
        "    parser.add_argument('--is_zip', default=1, choices=[0, 1], help='zip or not zip')\n",
        "\n",
        "    opt = parser.parse_known_args()[0]\n",
        "    return opt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 训练速度模型"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch\n",
        "from torch.utils.data import random_split,DataLoader\n",
        "from time import time\n",
        "from tqdm import tqdm\n",
        "import numpy as np\n",
        "import random\n",
        "import argparse\n",
        "import csv\n",
        "\n",
        "def set_seed(seed=1225):\n",
        "    \n",
        "    # 固定随机种子\n",
        "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
        "    random.seed(seed)\n",
        "    np.random.seed(seed)\n",
        "    torch.manual_seed(seed)\n",
        "    torch.cuda.manual_seed(seed)\n",
        "    torch.backends.cudnn.deterministic = True\n",
        "    torch.backends.cudnn.benchmark = False\n",
        "    msg =f\"\"\"\n",
        "    =============================================\n",
        "        seed = {seed}\n",
        "        torch.backends.cudnn.deterministic = True\n",
        "        torch.backends.cudnn.benchmark = False\n",
        "    =============================================\n",
        "    \"\"\"\n",
        "    print(msg)\n",
        "    \n",
        "def judge_file(folder_path='runs_trackA'):\n",
        "    folder_path = folder_path\n",
        "\n",
        "    if os.path.exists(folder_path) and os.path.isdir(folder_path):\n",
        "        print(\"文件夹存在\")\n",
        "    else:\n",
        "        os.system(f\"mkdir -p {folder_path}\")\n",
        "\n",
        "def train(opt):\n",
        "\n",
        "    log_p = opt.log_dir\n",
        "    # 设置随机种子\n",
        "    set_seed(opt.seed)\n",
        "    train_split = opt.train_ratio\n",
        "    train_batch = opt.batch\n",
        "    val_batch = opt.val_batch\n",
        "    num_epoch = opt.epochs\n",
        "    n_layers = opt.n_layers\n",
        "    n_hidden = opt.n_hidden\n",
        "    space_dim = opt.space_dim\n",
        "    mlp_ratio = opt.mlp_ratio\n",
        "    fun_dim = opt.fun_dim\n",
        "    out_dim = opt.out_dim\n",
        "    lr = opt.lr\n",
        "    max_lr = opt.max_lr\n",
        "    act = opt.act\n",
        "    slice_num = opt.slice_num\n",
        "    n_head = opt.n_head\n",
        "    \n",
        "    # 创建保存权重的文件夹\n",
        "    save_file = 'runs_trackCIKM'\n",
        "    judge_file(save_file)\n",
        "    log_file = os.path.join(save_file, log_p)\n",
        "    judge_file(log_file)\n",
        "\n",
        "    # 保存训练过程到csv文件中\n",
        "    f = open(f'{log_file}/train_process.csv', mode='w')\n",
        "    fc = csv.writer(f)\n",
        "    fc.writerow([\"train_loss\", \"val_loss\", 'lr'])\n",
        "\n",
        "    with open(f'{log_file}/args.txt', mode='w') as fargs:\n",
        "        fargs.write(str(opt))\n",
        "    \n",
        "    print(f\"saving args to {log_file}\")\n",
        "    \n",
        "    \n",
        "    \n",
        "    # 1. -加载CIKM的 速度 数据集\n",
        "    vel_paths, mesh_paths = read_dir()\n",
        "    dataset = VelDataset(vel_paths)\n",
        "    \n",
        "\n",
        "    # 检查gpu是否可用\n",
        "    if torch.cuda.is_available():\n",
        "        device = torch.device(\"cuda\")\n",
        "        print(\"GPU is available\")\n",
        "    else:\n",
        "        device = torch.device(\"cpu\")\n",
        "        print(\"GPU is not available, using CPU\")\n",
        "\n",
        "\n",
        "    # 2. -划分数据集为训练集和验证集 0.8 : 0.2\n",
        "    # 可以试一下 0.9:0.1 可能效果会好很多\n",
        "    train_size = int(train_split * len(dataset))\n",
        "    val_size = len(dataset) - train_size\n",
        "    train_dataset, val_dataset = random_split(dataset, [train_size, val_size])\n",
        "\n",
        "\n",
        "    # 3. -创建数据集加载器\n",
        "    train_loader = DataLoader(train_dataset, batch_size=train_batch, shuffle=True)\n",
        "    val_loader = DataLoader(val_dataset, batch_size=val_batch, shuffle=False)\n",
        "\n",
        "\n",
        "    # 4. -实例化模型\n",
        "    model = Model(\n",
        "        n_layers=n_layers, \n",
        "        n_hidden=n_hidden, \n",
        "        space_dim=space_dim, \n",
        "        mlp_ratio=mlp_ratio, \n",
        "        fun_dim=fun_dim,\n",
        "        act=act,\n",
        "        n_head=n_head,\n",
        "        slice_num=slice_num,\n",
        "        out_dim=out_dim\n",
        "        )\n",
        "\n",
        "    \n",
        "        \n",
        "    \n",
        "    model.to(device)\n",
        "\n",
        "    params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
        "    print(f\"---模型的参数为：{params} ---\")\n",
        "    \n",
        "\n",
        "    # 5. -定义损失函数\n",
        "    criterion = LpLoss(size_average=True)\n",
        "\n",
        "    # 6. -定义优化器 修改不同的参数，就不能加载预训练权重，会重新开始训练，最好的方法就是修改存储weights的文件夹的名称 学习率不能太大，太难了乱跑 学习率需要小一些 下降的速度也需要慢一点，太快了容易提前收敛了\n",
        "    #  之后需要修改lr lr_decay,获取弄的比较好的效果\n",
        "    optimizer = optim.Adam(model.parameters(), lr=0.0001)\n",
        "    # optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.00005, betas=(0.9, 0.999))\n",
        "    \n",
        "    # 7. -开始训练\n",
        "    best_val_loss = 1000\n",
        "    best_train_loss = 1000\n",
        "    # 把学习率下降调小\n",
        "\n",
        "    scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
        "        optimizer,\n",
        "        max_lr=max_lr,\n",
        "        epochs=num_epoch,\n",
        "        div_factor=1e4, pct_start=0.25, final_div_factor=1e4,\n",
        "        steps_per_epoch=len(train_loader)\n",
        "    )\n",
        "    \n",
        "    # scheduler = optim.lr_scheduler.CyclicLR(optimizer, base_lr=0.000001, max_lr=0.0008, step_size_up=800, mode='triangular')\n",
        "    grad_clip = 1000.0\n",
        "    for epoch in range(num_epoch):\n",
        "        t1 = time()\n",
        "\n",
        "        model.train()\n",
        "        train_loss = 0\n",
        "        for X_batch, y_batch in tqdm(train_loader):\n",
        "            X_batch, y_batch = X_batch.to(device), y_batch.to(device)\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "            outputs = model(X_batch)\n",
        "\n",
        "            loss = criterion(outputs,y_batch)\n",
        "    \n",
        "            loss.backward()     \n",
        "            nn.utils.clip_grad_norm_(model.parameters(), grad_clip)\n",
        "            optimizer.step()\n",
        "            scheduler.step()\n",
        "                  \n",
        "            train_loss += loss.item()\n",
        "            # break\n",
        "        train_loss /= len(train_loader)\n",
        "\n",
        "        model.eval()\n",
        "        val_loss = 0\n",
        "        with torch.no_grad():\n",
        "            for X_batch, y_batch in val_loader:\n",
        "                X_batch, y_batch = X_batch.to(device), y_batch.to(device)\n",
        "                outputs = model(X_batch)\n",
        "                # y_batch = y_batch.reshape((-1,3))\n",
        "                # outputs = outputs.reshape((-1,3))\n",
        "                loss = criterion(outputs, y_batch)\n",
        "                val_loss += loss.item()\n",
        "        val_loss /= len(val_loader)\n",
        "\n",
        "        # 8. -保存模型最后一轮权重\n",
        "        last_weights_path = f'{log_file}/last_weights.pt'\n",
        "        torch.save(model.state_dict(),last_weights_path)\n",
        "\n",
        "        # 9. -保存最优权重\n",
        "        if val_loss <= best_val_loss and train_loss <= best_train_loss:\n",
        "            best_val_loss = val_loss\n",
        "            best_train_loss = train_loss\n",
        "            best_weight_path = f'{log_file}/best_weights.pt'\n",
        "            torch.save(model.state_dict(), best_weight_path)\n",
        "        \n",
        "        current_lr = optimizer.param_groups[0]['lr']\n",
        "\n",
        "        \n",
        "        t2 =time()\n",
        "        t = t2 - t1\n",
        "        print(f'Epoch {epoch+1}/{num_epoch}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, Current_lr: {current_lr}, Spend Time: {round(t,2)}s')\n",
        "        fc.writerow([train_loss, val_loss, current_lr])\n",
        "\n",
        "    f.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "from time import sleep\n",
        "from tqdm import tqdm\n",
        "import csv\n",
        "def mmmodel():\n",
        "    f = open('submission/Answer.csv', mode='w')\n",
        "    fc = csv.writer(f)\n",
        "    fc.writerow(['', 'Cd'])\n",
        "    for i in tqdm(range(50)):\n",
        "        sleep(1)\n",
        "        fc.writerow([i, 0.415])\n",
        "    f.close()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 测试速度代码"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [],
      "source": [
        "def test(opt, train=False):\n",
        "    log = opt.log_dir\n",
        "    mode = opt.result\n",
        "    output_file = opt.output_file\n",
        "    # 检查gpu是否可用\n",
        "    if torch.cuda.is_available():\n",
        "        device = torch.device(\"cuda\")\n",
        "        print(\"GPU is available\")\n",
        "    else:\n",
        "        device = torch.device(\"cpu\")\n",
        "        print(\"GPU is not available, using CPU\")\n",
        "\n",
        "    # 加载模型，进行测试\n",
        "    n_layers = opt.n_layers\n",
        "    n_hidden = opt.n_hidden\n",
        "    space_dim = opt.space_dim\n",
        "    mlp_ratio = opt.mlp_ratio\n",
        "    fun_dim = opt.fun_dim\n",
        "    act = opt.act\n",
        "    slice_num = opt.slice_num\n",
        "    n_head = opt.n_head\n",
        "    out_dim = opt.out_dim\n",
        "    model = Model(\n",
        "        n_layers=n_layers, \n",
        "        n_hidden=n_hidden, \n",
        "        space_dim=space_dim, \n",
        "        mlp_ratio=mlp_ratio, \n",
        "        fun_dim=fun_dim,\n",
        "        act=act,\n",
        "        n_head=n_head,\n",
        "        slice_num=slice_num,\n",
        "        out_dim=out_dim\n",
        "        )\n",
        "    \n",
        "    \n",
        "    os.system(f\"mkdir -p runs_trackCIKM/{log}/{output_file}\")\n",
        "\n",
        "    # 加载最优的权重模型\n",
        "    if train:\n",
        "        best_weights_path = f'runs_trackCIKM/{log}/best_weights.pt'\n",
        "        model.load_state_dict(torch.load(best_weights_path))\n",
        "    else:\n",
        "        best_weights_path = f'transolver_vel.pt'\n",
        "        model.load_state_dict(torch.load(best_weights_path))\n",
        "\n",
        "     \n",
        "    # 加载到GPU上\n",
        "    model.to(device)\n",
        "    \n",
        "    model.eval()\n",
        "    # ----CIKM测试集-----------\n",
        "    vel_paths, mesh_paths = read_dir(test=True)\n",
        "    dataset = VelTestDataset(vel_paths)\n",
        "    \n",
        "    test_loader = DataLoader(dataset, batch_size=1, shuffle=False)\n",
        "\n",
        "    # 进行测试\n",
        "    for X_batch, p in test_loader:\n",
        "\n",
        "        # 将数据部署到GPU上\n",
        "        X_batch= X_batch.to(device)\n",
        "\n",
        "        outputs = model(X_batch)\n",
        "\n",
        "        # 将输出到cpu上\n",
        "        outputs = outputs.to('cpu')\n",
        "\n",
        "        outputs = outputs.detach().numpy().reshape((-1,3)).astype(np.float32)\n",
        "\n",
        "        # print(np.min(outputs))\n",
        "\n",
        "        # 处理保存数据的路径 \n",
        "        number = str(p).split(\".\")[0][-3:]\n",
        "        path = f'{output_file}/vel_{number}.npy'\n",
        "\n",
        "        # 保存数据\n",
        "        np.save(path,outputs)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 先测试最优权重"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "if __name__ == \"__main__\":\n",
        "    opt = get_args()\n",
        "    test(opt)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 训练并验证最优代码"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "if __name__ == \"__main__\":\n",
        "\n",
        "    opt = get_args()\n",
        "    train(opt=opt)\n",
        "    test(opt,train=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y9uq00gkYq1B"
      },
      "source": [
        "## Dataset2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### 训练obj模型,直接导出文件保存"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NjbtsOF5YumS",
        "outputId": "4952429c-6858-481e-df56-be9a0469cd14"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 50/50 [00:00<00:00, 49648.48it/s]\n"
          ]
        }
      ],
      "source": [
        "\n",
        "result = mmmodel()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7YRUsrbjZ-pm"
      },
      "source": [
        "# 结果保存\n",
        "（结果在此模块下压缩成submission.zip，提交结果不可以通过外部链接下载）"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zPkn9fEPlg4y"
      },
      "outputs": [],
      "source": [
        "import zipfile\n",
        "with zipfile.ZipFile('submission.zip','w') as zipObj:\n",
        "    file_list = os.listdir('submission')\n",
        "    for f in file_list:\n",
        "        p = os.path.join('submission/', f)\n",
        "        zipObj.write(p)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "Q0mY71_DZwFx",
        "YQE2mV-JZyCo"
      ],
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
